import os
import shutil
import subprocess
import sys


def move_images_and_labelme(src_dir, dst_dir, form):
    if not os.path.exists(dst_dir):
        os.makedirs(dst_dir)
    for root, dirs, files in os.walk(src_dir):
        for file in files:
            if file.endswith(form) or file.endswith('.json'):
                src_file = os.path.join(root, file)
                dst_file = os.path.join(dst_dir, file)

                shutil.copy2(src_file, dst_file)
                print(f"cp {src_file} to {dst_file}")


def main(args):
    # create root directory to store these pys output
    root_directory = f"/home/jacy/GDW/data_{args[3]}"
    if not os.path.exists(root_directory):
        os.makedirs(root_directory)
    else:
        print(f"{root_directory} is already existing")

    # corresponding to 0.py
    source_dir = args[1]
    image_form = args[2]
    data_name_string = args[3]
    target_dir = f'/home/jacy/GDW/data_{data_name_string}/labelme'
    move_images_and_labelme(source_dir, target_dir, image_form)

    # corresponding to 1,2,3,4,5.py
    scripts = [
        ("1_generate_coco.py", data_name_string),
        ("2_generate_cfg.py", data_name_string),
        ("3_train.py", data_name_string),
        ("4_test.py", data_name_string),
        ("5_cook.py", data_name_string),
    ]

    current_dir = '.'
    for script_name, arg in scripts:
        command = ["python", f"{current_dir}/{script_name}", arg]

        try:
            subprocess.run(command, check=True)
            print(f"\033[32m Executed {script_name} with argument {arg} successfully. \033[0m")

        except subprocess.CalledProcessError as e:
            print(f"\033[31m Error executing {script_name} with argument {arg}: \033[0m")
            print(e)
            break

    print(f"\033[32m All scripts have been attempted to execute. \033[0m")


if __name__ == '__main__':
    # usage: python all_in_one.py source_data_dir data_format data_source_abbreviation
    # example: python all_in_one.py /media/share-rw/GOODWE-POC/NG1/Jan/shaixuan/Chip/Resistor/6332-OK_LY .jpg test
    main(sys.argv)
